import cv2
import numpy as np
import matplotlib.pyplot as plt

def auto_correct_and_detect_water_level(image_path):
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 预处理：高斯滤波 + 二值化
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    _, binary = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    cv2.imshow('rest', blurred)
    # 膨胀关闭操作，去除瓶间缝隙
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (7, 7))
    closed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
    cv2.imshow('reslt', closed)

    # 找瓶子轮廓
    contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    bottle_contours = [c for c in contours if cv2.contourArea(c) > 1000]

    results_img = img.copy()
    for contour in bottle_contours:
        x, y, w, h = cv2.boundingRect(contour)
        bottle = gray[y:y+h, x:x+w]

        # 分析瓶子中间部分做水平投影
        roi = bottle[int(h*0.2):int(h*0.8), :]
        projection = np.sum(255 - roi, axis=1)  # 水平投影
        # 修复：转换为 float32 类型
        projection_smooth = cv2.GaussianBlur(projection.astype(np.float32).reshape(-1, 1), (5, 1), 0).flatten()


        # 寻找水位线（最小投影值的位置）
        min_index = np.argmin(projection_smooth)
        water_level_y = int(h * 0.2) + min_index

        # 判断是否合格：若水位高于瓶子高度的中点则合格
        is_ok = water_level_y < y + h*0.2
        label = "OK" if is_ok else "NG"
        color = (0, 255, 0) if is_ok else (0, 0, 255)

        cv2.rectangle(results_img, (x, y), (x+w, y+h), color, 2)
        cv2.putText(results_img, label, (x, y+30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
        cv2.line(results_img, (x, water_level_y), (x+w, water_level_y), color, 2)

    return results_img

# 运行检测
image_path = "10/6.png"
result_image = auto_correct_and_detect_water_level(image_path)

cv2.imshow('result', result_image)
cv2.waitKey(0)
cv2.destroyAllWindows()
